﻿Imports System.Security.Cryptography
Imports System.IO
Imports System.Runtime.Serialization
Imports System.Runtime.Serialization.Formatters.Binary
Imports System.Runtime.CompilerServices

''' <summary>
''' Mistery message
''' </summary>
''' <remarks></remarks>
<Serializable()>
Public Class WebMessage
    ''' <summary>
    ''' clear
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property Sender As String

    ''' <summary>
    ''' encrypted with receiver public
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property Token As Byte()

    ''' <summary>
    ''' encrypted with receiver public
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property MessageID As Byte()

    ''' <summary>
    ''' Hash of MessageID encrypted with sender private 
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property MessageIDSignaure As Byte()

    ''' <summary>
    ''' encrypted with receiver public
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property Key() As Byte()
    ''' <summary>
    ''' encrypted with receiver public
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property IV As Byte()

    ''' <summary>
    ''' encrypted with receiver public, contains the serilizied message input
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property Message As Byte()

    ''' <summary>
    ''' Hash of message encrypted with sender private 
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Property Signature As Byte()


    Public Shared Function GenerateMessage(Sender As String, message As Object,
                                           receiver_rsa As RSACryptoServiceProvider,
                                           token As String) As WebMessage

        Dim formatter As IFormatter = New BinaryFormatter()

        Dim RMCrypto As New RijndaelManaged()
        RMCrypto.GenerateIV()
        RMCrypto.GenerateKey()

        Dim ms As New MemoryStream()

        formatter.Serialize(ms, message)
        Dim result_message() As Byte = ms.ToArray

        ms = New MemoryStream()
        Dim CryptStream As New CryptoStream(ms, RMCrypto.CreateEncryptor(RMCrypto.Key, RMCrypto.IV), CryptoStreamMode.Write)
        CryptStream.Write(result_message, 0, result_message.Length)
        CryptStream.FlushFinalBlock()

        Dim message_id_bytes As Byte() = Guid.NewGuid.ToString.getBytes

        Dim my_rsa As RSACryptoServiceProvider = CryptStream.getGlobalObject(Of IMyRsaProvider)().getMySra

        Dim answer As New WebMessage With {.Sender = Sender,
                                           .IV = receiver_rsa.Encrypt(RMCrypto.IV, False),
                                           .Key = receiver_rsa.Encrypt(RMCrypto.Key, False),
                                           .Message = ms.ToArray,
                                           .Signature = my_rsa.SignData(result_message, New SHA1CryptoServiceProvider),
                                           .Token = receiver_rsa.Encrypt(token.getBytes, False),
                                           .MessageID = receiver_rsa.Encrypt(message_id_bytes, False),
                                           .MessageIDSignaure = my_rsa.SignData(message_id_bytes, New SHA1CryptoServiceProvider)
                                           }
        Return answer
    End Function

    Private Shared message_ids As New HashSet(Of String)
    Private Const replay_store_size As Integer = 100000

    Public Shared Function readMessage(message As WebMessage, sender_rsa As RSACryptoServiceProvider) As Object
        If message Is Nothing Then Return Nothing
        Dim signature() As Byte = message.Signature

        Dim my_rsa As RSACryptoServiceProvider = message.getGlobalObject(Of IMyRsaProvider)().getMySra

        'checking unique id
        Dim message_id_bytes As Byte() = my_rsa.Decrypt(message.MessageID, False)
        If Not sender_rsa.VerifyData(message_id_bytes, New SHA1CryptoServiceProvider, message.MessageIDSignaure) Then
            Return Nothing
        End If
        Dim message_id As String = message_id_bytes.getString
        If message_ids.Contains(message_id) Then Return Nothing 'replay attack!!!
        If message_ids.Count > replay_store_size Then
            message_ids = New HashSet(Of String)
        End If
        message_ids.Add(message_id)

        Dim Key As Byte() = my_rsa.Decrypt(message.Key, False)
        Dim IV As Byte() = my_rsa.Decrypt(message.IV, False)

        'Create a new instance of the RijndaelManaged class
        'and decrypt the stream.
        Dim RMCrypto As New RijndaelManaged()


        'Create an instance of the CryptoStream class, pass it the NetworkStream, and decrypt 
        'it with the Rijndael class using the key and IV.
        Dim ms As New MemoryStream(message.Message)
        Dim CryptStream As New CryptoStream(ms, RMCrypto.CreateDecryptor(Key, IV), CryptoStreamMode.Read)


        Dim output As New MemoryStream
        Dim buffer(1024 - 1) As Byte
        Dim read As Integer = CryptStream.Read(buffer, 0, buffer.Length)
        While read > 0
            output.Write(buffer, 0, read)
            read = CryptStream.Read(buffer, 0, buffer.Length)
        End While
        CryptStream.Flush()
        CryptStream.Dispose()
        ms.Dispose()

        Dim received_message() As Byte = output.ToArray

        If Not sender_rsa.VerifyData(received_message, New SHA1CryptoServiceProvider, signature) Then
            Return Nothing
        End If

        'ok message verificated and descripted!
        Dim formatter As IFormatter = New BinaryFormatter()
        ms = New MemoryStream(received_message)
        Dim message_object As Object = formatter.Deserialize(ms)
        ms.Dispose()

        Return message_object

    End Function

End Class


<GlobalAvalibleObject()>
Public Interface IMyRsaProvider
    Function getMySra() As RSACryptoServiceProvider
End Interface

<GlobalAvalibleObjectImplementation(implementation_of:=GetType(IMyRsaProvider), overrides_exsisting:=False, singleton:=True)>
Public Class DefaultMyRsaProvider
    Implements IMyRsaProvider

    Public Function getMySra() As RSACryptoServiceProvider Implements IMyRsaProvider.getMySra
        ' Create the CspParameters object and set the key container 
        '  name used to store the RSA key pair.
        Dim cp As New CspParameters()
        cp.KeyContainerName = "MysteryKey"

        ' Create a new instance of RSACryptoServiceProvider that accesses
        ' the key container MyKeyContainerName.
        Dim rsa As New RSACryptoServiceProvider(cp)

        Return rsa
    End Function
End Class

Public Module RsaExntesions
    ''' <summary>
    ''' read a string encrypted with my public key
    ''' </summary>
    ''' <param name="encrypted"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    <Extension()>
    Public Function readAEncryptedString(encrypted As Byte()) As String
        Return encrypted.getGlobalObject(Of IMyRsaProvider)().getMySra.Decrypt(encrypted, False).getString
    End Function
End Module
